from django.core.management.base import BaseCommand
from django.db import transaction
from movies.models import Movie, Genre
import requests
import os
from tqdm import tqdm

class Command(BaseCommand):
    help = '从TMDB API导入电影'

    def add_arguments(self, parser):
        parser.add_argument('--count', type=int, default=100, help='要导入的电影数量')
        parser.add_argument('--api_key', type=str, required=True, help='TMDB API密钥')

    def handle(self, *args, **options):
        count = options['count']
        api_key = options['api_key']
        
        # 获取流行电影
        self.import_popular_movies(api_key, count)
        
        self.stdout.write(self.style.SUCCESS('导入完成!'))
    
    def import_popular_movies(self, api_key, count):
        """从TMDB导入流行电影"""
        # 获取所有现有的类型
        existing_genres = {genre.name: genre for genre in Genre.objects.all()}
        
        # 获取TMDB的电影类型列表
        genres_url = f"https://api.themoviedb.org/3/genre/movie/list?api_key={api_key}&language=zh-CN"
        genres_response = requests.get(genres_url)
        genres_data = genres_response.json()
        
        # 创建类型ID到名称的映射
        genre_mapping = {}
        for genre in genres_data.get('genres', []):
            genre_mapping[genre['id']] = genre['name']
            if genre['name'] not in existing_genres:
                new_genre = Genre.objects.create(name=genre['name'])
                existing_genres[genre['name']] = new_genre
        
        # 分页获取流行电影
        movies_added = 0
        page = 1
        
        with tqdm(total=count) as pbar:
            while movies_added < count:
                movies_url = f"https://api.themoviedb.org/3/movie/popular?api_key={api_key}&language=zh-CN&page={page}"
                movies_response = requests.get(movies_url)
                movies_data = movies_response.json()
                
                if 'results' not in movies_data:
                    break
                
                for movie_data in movies_data['results']:
                    if movies_added >= count:
                        break
                    
                    title = movie_data['title']
                    release_date = movie_data.get('release_date', '')
                    year = int(release_date[:4]) if release_date and len(release_date) >= 4 else 2000
                    
                    # 检查电影是否已存在
                    if Movie.objects.filter(title=title, year=year).exists():
                        continue
                    
                    # 获取更详细的电影信息
                    movie_details_url = f"https://api.themoviedb.org/3/movie/{movie_data['id']}?api_key={api_key}&language=zh-CN"
                    movie_details_response = requests.get(movie_details_url)
                    movie_details = movie_details_response.json()
                    
                    # 创建电影
                    with transaction.atomic():
                        movie = Movie.objects.create(
                            title=title,
                            year=year,
                            description=movie_data.get('overview', '') or movie_details.get('overview', '暂无简介'),
                            poster_url=f"https://image.tmdb.org/t/p/w500{movie_data.get('poster_path', '')}" if movie_data.get('poster_path') else None
                        )
                        
                        # 添加类型
                        for genre_id in movie_data.get('genre_ids', []):
                            if genre_id in genre_mapping:
                                genre_name = genre_mapping[genre_id]
                                genre = existing_genres[genre_name]
                                movie.genres.add(genre)
                        
                        self.stdout.write(f'导入电影: {title} ({year})')
                        movies_added += 1
                        pbar.update(1)
                
                page += 1
                
                # 如果没有更多页面，跳出循环
                if page > movies_data.get('total_pages', 1):
                    break 